import os
import platform
import signal
import readline
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
 
 
torch.cuda.empty_cache()

# 在启动的py文件里添加一段组织模型路径的代码
# 模型位置使用绝对路径并输出一下，可以检查一下对不对
current_path = os.path.abspath(__file__) # 当前脚本文件绝对路径
current_dir = os.path.dirname(current_path) # 当前脚本所在文件夹绝对路径
model_path = os.path.join(current_dir, ".", "baichuan-7B") # 拼接路径
print("模型路径：", model_path) 

# 修改 AutoTokenizer.from_pretrained 和 AutoModel.from_pretrained 的第一个参数为前面的 model_path
# AutoModel.from_pretrained(...).cuda() 改为 .float() 使用cpu运行。
# tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
# model = AutoModel.from_pretrained(model_path, trust_remote_code=True).float()  

# 加载模型
tokenizer = AutoTokenizer.from_pretrained("./baichuan-7B", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("./baichuan-7B", device_map="auto", trust_remote_code=True)
 
# 多显卡支持，使用下面两行代替上面一行，将num_gpus改为你实际的显卡数量
# from utils import load_model_on_gpus
# model = load_model_on_gpus("THUDM/chatglm2-6b", num_gpus=2)
model = model.eval()

os.environ['CUDA_VISIBLE_DEVICES']='0'
os_name = platform.system()
clear_command = 'cls' if os_name == 'Windows' else 'clear'
stop_stream = False


def build_prompt(history):
    prompt = "欢迎使用 baichuan-7B 模型，输入内容即可进行对话，clear 清空对话历史，stop 终止程序"
    for query, response in history:
        prompt += f"\n\n用户：{query}"
        prompt += f"\n\n baichuan-7B：{response}"
    return prompt


def signal_handler(signal, frame):
    global stop_stream
    stop_stream = True


def main():
    past_key_values, history = None, []
    global stop_stream
    print("欢迎使用 baichuan-7B 模型，输入内容即可进行对话，clear 清空对话历史，stop 终止程序")
    while True:
	    # torch.cuda.empty_cache()
        question = input("\n用户：")
        if question.strip() == "stop":
            break
        if question.strip() == "clear":
            past_key_values, history = None, []
            os.system(clear_command)
            print("欢迎使用 baichuan-7B 模型，输入内容即可进行对话，clear 清空对话历史，stop 终止程序")
            continue
        print("\nbaichuan-7B：", end="")
        current_length = 0
        # 接收 POST 请求的数据
         
        question += '->'

        inputs = tokenizer(question, return_tensors='pt')
        inputs = inputs.to('cuda:0')
        pred = model.generate(**inputs, max_new_tokens=1024, repetition_penalty=1.1)

        text = tokenizer.decode(pred.cpu()[0], skip_special_tokens=True)
        print("", text) 
        print("")


if __name__ == "__main__":
    main()
